import torch
import numpy as np
import os

from matplotlib import pyplot as plt

from model_advanced import HybridWorldModel
from advanced_features import create_advanced_state_representation
from advanced_training import advanced_train
from ensemble_prediction import EnsemblePrediction
from utils import prepare_batch
from torch.utils.data import DataLoader, TensorDataset
from config import *
from environment import MazeEnvironment


def generate_advanced_dataset(num_samples=NUM_SAMPLES, seed=RANDOM_SEED):
    """生成高级数据集，包含增强状态表示"""
    # 创建环境
    env = MazeEnvironment(seed=seed)

    # 准备初始布局和静态掩码
    initial_layout = env.get_state().copy()
    np.save('../initial_layout.npy', initial_layout)

    static_mask = np.zeros_like(initial_layout, dtype=bool)
    static_mask[initial_layout == OBSTACLE] = True
    static_mask[initial_layout == START] = True
    static_mask[initial_layout == END] = True
    static_mask[initial_layout == COIN] = True
    np.save('../static_mask.npy', static_mask)

    # 存储数据
    states = []
    actions = []
    next_states = []
    agent_positions = []
    agent_movements = []

    # 记录统计信息
    movement_samples = 0
    total_samples = 0

    # 生成数据
    while len(states) < num_samples:
        # 定期重置环境
        if len(states) % 50 == 0:
            env.reset_without_changing_layout()

        # 获取当前状态
        current_state = env.get_state().copy()
        current_agent_pos = env.get_agent_position(current_state)

        if current_agent_pos is None:
            env.reset_without_changing_layout()
            continue

        # 找出可用动作
        valid_actions = []
        move_actions = []

        for action in range(4):  # UP, DOWN, LEFT, RIGHT
            next_state = env.get_next_state_after_action(action)
            next_agent_pos = env.get_agent_position(next_state)

            if next_agent_pos and next_agent_pos != current_agent_pos:
                move_actions.append(action)

            if next_agent_pos:
                valid_actions.append(action)

        # 选择动作，优先选择移动动作
        if move_actions and (movement_samples < AGENT_MOVEMENT_RATIO * total_samples or np.random.random() < 0.9):
            action = np.random.choice(move_actions)
            is_movement = True
        elif valid_actions:
            action = np.random.choice(valid_actions)
            is_movement = (action in move_actions)
        else:
            action = np.random.randint(0, 4)
            is_movement = False

        # 获取下一个状态
        next_state = env.get_next_state_after_action(action)
        next_agent_pos = env.get_agent_position(next_state)

        # 只有当能找到智能体位置时才添加样本
        if current_agent_pos and next_agent_pos:
            # 计算相对移动
            dx = next_agent_pos[0] - current_agent_pos[0]
            dy = next_agent_pos[1] - current_agent_pos[1]

            # 创建增强状态表示
            current_advanced = create_advanced_state_representation(current_state)
            next_advanced = create_advanced_state_representation(next_state)

            # 添加样本
            states.append(current_advanced.flatten())
            actions.append(action)
            next_states.append(next_advanced.flatten())
            agent_positions.append([next_agent_pos[0], next_agent_pos[1]])
            agent_movements.append([dx, dy, next_agent_pos[0], next_agent_pos[1]])

            # 更新计数器
            total_samples += 1
            if is_movement:
                movement_samples += 1

            # 执行动作
            env.step(action)
        else:
            # 如果找不到智能体，重置环境
            env.reset_without_changing_layout()

    # 确保样本数量正确
    states = states[:num_samples]
    actions = actions[:num_samples]
    next_states = next_states[:num_samples]
    agent_positions = agent_positions[:num_samples]
    agent_movements = agent_movements[:num_samples]

    # 输出移动样本比例
    movement_ratio = movement_samples / total_samples
    print(f"生成的高级数据集中包含智能体移动的样本比例: {movement_ratio:.2f}")

    # 转换为numpy数组
    states = np.array(states)
    actions = np.array(actions)
    next_states = np.array(next_states)
    agent_positions = np.array(agent_positions)
    agent_movements = np.array(agent_movements)

    # 随机打乱数据
    rng = np.random.RandomState(seed)
    indices = rng.permutation(num_samples)
    states = states[indices]
    actions = actions[indices]
    next_states = next_states[indices]
    agent_positions = agent_positions[indices]
    agent_movements = agent_movements[indices]

    # 分割为训练、验证和测试集
    num_train = int(num_samples * TRAIN_SPLIT)
    num_val = int(num_samples * VAL_SPLIT)

    train_indices = np.arange(num_train)
    val_indices = np.arange(num_train, num_train + num_val)
    test_indices = np.arange(num_train + num_val, num_samples)

    train_data = (states[train_indices], actions[train_indices], next_states[train_indices],
                  agent_positions[train_indices], agent_movements[train_indices])
    val_data = (states[val_indices], actions[val_indices], next_states[val_indices],
                agent_positions[val_indices], agent_movements[val_indices])
    test_data = (states[test_indices], actions[test_indices], next_states[test_indices],
                 agent_positions[test_indices], agent_movements[test_indices])

    return train_data, val_data, test_data


def prepare_datasets(train_data, val_data, test_data):
    """准备数据集和数据加载器"""
    # 解包数据
    train_states, train_actions, train_next_states, train_agent_positions, train_agent_movements = train_data
    val_states, val_actions, val_next_states, val_agent_positions, val_agent_movements = val_data
    test_states, test_actions, test_next_states, test_agent_positions, test_agent_movements = test_data

    # 创建张量
    train_states_tensor = torch.FloatTensor(train_states)
    train_actions_tensor = torch.LongTensor(train_actions)
    train_next_states_tensor = torch.FloatTensor(train_next_states)
    train_agent_positions_tensor = torch.FloatTensor(train_agent_positions)
    train_agent_movements_tensor = torch.FloatTensor(train_agent_movements)

    val_states_tensor = torch.FloatTensor(val_states)
    val_actions_tensor = torch.LongTensor(val_actions)
    val_next_states_tensor = torch.FloatTensor(val_next_states)
    val_agent_positions_tensor = torch.FloatTensor(val_agent_positions)
    val_agent_movements_tensor = torch.FloatTensor(val_agent_movements)

    test_states_tensor = torch.FloatTensor(test_states)
    test_actions_tensor = torch.LongTensor(test_actions)
    test_next_states_tensor = torch.FloatTensor(test_next_states)
    test_agent_positions_tensor = torch.FloatTensor(test_agent_positions)
    test_agent_movements_tensor = torch.FloatTensor(test_agent_movements)

    # 创建数据集
    train_dataset = TensorDataset(
        train_states_tensor, train_actions_tensor, train_next_states_tensor,
        train_agent_positions_tensor, train_agent_movements_tensor
    )

    val_dataset = TensorDataset(
        val_states_tensor, val_actions_tensor, val_next_states_tensor,
        val_agent_positions_tensor, val_agent_movements_tensor
    )

    test_dataset = TensorDataset(
        test_states_tensor, test_actions_tensor, test_next_states_tensor,
        test_agent_positions_tensor, test_agent_movements_tensor
    )

    return train_dataset, val_dataset, test_dataset


def evaluate_advanced_model(model, test_dataset, device=None):
    """评估高级混合模型的性能"""
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model.eval()
    test_loader = DataLoader(test_dataset, batch_size=64)

    # 初始化指标
    env_acc = 0
    agent_abs_acc = 0
    agent_rel_acc = 0
    total_samples = 0

    with torch.no_grad():
        for states, actions, next_states, agent_positions, agent_movements in test_loader:
            states = states.to(device)
            next_states = next_states.to(device)
            agent_positions = agent_positions.to(device)
            agent_movements = agent_movements.to(device)

            # 前向传递
            env_outputs, agent_outputs = model(states)

            # 计算环境准确率
            env_acc_batch = calculate_env_accuracy(env_outputs, next_states)

            # 计算智能体绝对位置准确率
            abs_pred = torch.round(agent_outputs[:, 2:]).long()
            abs_true = agent_positions.long()
            abs_acc_batch = torch.mean((abs_pred == abs_true).all(dim=1).float()).item()

            # 计算智能体相对移动准确率
            rel_pred = torch.round(agent_outputs[:, :2]).long()
            rel_true = agent_movements[:, :2].long()
            rel_acc_batch = torch.mean((rel_pred == rel_true).all(dim=1).float()).item()

            # 累加指标
            batch_size = states.size(0)
            env_acc += env_acc_batch * batch_size
            agent_abs_acc += abs_acc_batch * batch_size
            agent_rel_acc += rel_acc_batch * batch_size
            total_samples += batch_size

    # 计算平均指标
    env_acc /= total_samples
    agent_abs_acc /= total_samples
    agent_rel_acc /= total_samples

    print("======= 高级模型测试结果 =======")
    print(f"环境状态准确率: {env_acc:.4f}")
    print(f"智能体绝对位置准确率: {agent_abs_acc:.4f}")
    print(f"智能体相对移动准确率: {agent_rel_acc:.4f}")
    print("===============================")

    return env_acc, agent_abs_acc, agent_rel_acc


def calculate_env_accuracy(env_outputs, next_states):
    """计算环境状态预测准确率"""
    batch_size = env_outputs.size(0)

    # 提取基本特征
    features_per_cell = next_states.size(1) // (ENV_SIZE * ENV_SIZE)
    next_states_reshaped = next_states.view(batch_size, ENV_SIZE, ENV_SIZE, features_per_cell)
    next_states_basic = next_states_reshaped[:, :, :, :CELL_FEATURES]

    # 计算准确率
    env_acc = 0
    for i in range(batch_size):
        env_pred = torch.argmax(env_outputs[i], dim=-1).view(-1)
        env_true = torch.argmax(next_states_basic[i], dim=-1).view(-1)
        env_acc += torch.mean((env_pred == env_true).float()).item()

    return env_acc / batch_size


def run_ensemble_episodes(num_episodes=5, steps_per_episode=20):
    """运行使用集成预测的测试episodes"""
    # 创建集成预测器
    predictor = EnsemblePrediction(model_path="models/hybrid_world_model_best.pth")

    # 创建可视化目录
    os.makedirs("advanced_episodes", exist_ok=True)

    # 运行episodes
    total_agent_correct = 0
    total_steps = 0

    for episode in range(num_episodes):
        # 创建环境
        env = MazeEnvironment(seed=RANDOM_SEED)

        # 可视化初始状态
        initial_state = env.get_state()
        visualize_grid(
            initial_state,
            title=f"高级预测 Episode {episode + 1} - 初始状态",
            save_path=f"advanced_episodes/episode_{episode + 1}_step_0.png"
        )

        # 运行步骤
        episode_correct = 0

        for step in range(steps_per_episode):
            # 获取当前状态
            current_state = env.get_state().copy()

            # 创建增强状态表示
            state_representation = create_advanced_state_representation(current_state).flatten()

            # 选择随机动作
            action = np.random.randint(0, 4)  # UP, DOWN, LEFT, RIGHT

            # 获取真实下一个状态
            true_next_state = env.get_next_state_after_action(action).copy()

            # 使用集成预测下一个智能体位置
            agent_pos_pred = predictor.predict(state_representation)

            # 获取真实的智能体位置
            true_agent_pos = env.get_agent_position(true_next_state)

            # 创建预测状态
            pred_next_state = current_state.copy()
            # 移除当前智能体位置
            agent_positions = np.where(pred_next_state == AGENT)
            if len(agent_positions[0]) > 0:
                for ax, ay in zip(agent_positions[0], agent_positions[1]):
                    pred_next_state[ax, ay] = EMPTY

            # 设置预测的智能体位置
            i, j = agent_pos_pred
            if 0 <= i < ENV_SIZE and 0 <= j < ENV_SIZE and pred_next_state[i, j] != OBSTACLE:
                pred_next_state[i, j] = AGENT

            # 执行动作
            env.step(action)

            # 计算准确率
            correct = (true_agent_pos is not None and
                       agent_pos_pred[0] == true_agent_pos[0] and
                       agent_pos_pred[1] == true_agent_pos[1])

            if correct:
                episode_correct += 1
                total_agent_correct += 1

            total_steps += 1

            # 计算整体准确率
            overall_acc = np.mean(pred_next_state == true_next_state)

            print(f"Episode {episode + 1}, Step {step + 1}: "
                  f"Overall Acc: {overall_acc:.4f}, Agent Acc: {1.0 if correct else 0.0}")

            # 可视化预测
            visualize_prediction(
                true_next_state,
                pred_next_state,
                action,
                save_path=f"advanced_episodes/episode_{episode + 1}_step_{step + 1}.png"
            )

        # 输出每个episode的准确率
        print(f"Episode {episode + 1} 智能体预测准确率: {episode_correct / steps_per_episode:.2f}")

    # 输出总体准确率
    overall_accuracy = total_agent_correct / total_steps
    print(f"总体智能体预测准确率: {overall_accuracy:.4f}")

    return overall_accuracy


def visualize_grid(grid, title="Environment", save_path=None):
    """可视化环境网格"""
    # 导入这里避免循环导入
    import matplotlib.pyplot as plt

    colors = {
        EMPTY: 'white',
        START: 'green',
        AGENT: 'blue',
        COIN: 'yellow',
        OBSTACLE: 'black',
        END: 'red'
    }

    names = {
        EMPTY: 'Empty',
        START: 'Start',
        AGENT: 'Agent',
        COIN: 'Coin',
        OBSTACLE: 'Obstacle',
        END: 'End'
    }

    # 创建网格可视化
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_title(title)

    # 绘制网格
    height, width = grid.shape
    for i in range(height):
        for j in range(width):
            cell_value = grid[i, j]
            color = colors.get(cell_value, 'gray')
            rect = plt.Rectangle((j, height - i - 1), 1, 1, color=color, ec='black')
            ax.add_patch(rect)

            # 添加文本标签
            if cell_value not in [EMPTY, OBSTACLE]:
                ax.text(j + 0.5, height - i - 0.5, names.get(cell_value, 'Unknown'),
                        ha='center', va='center', fontsize=8)

    # 设置坐标轴
    ax.set_xlim(0, width)
    ax.set_ylim(0, height)
    ax.set_xticks(range(width))
    ax.set_yticks(range(height))
    ax.set_xticklabels(range(width))
    ax.set_yticklabels(range(height - 1, -1, -1))
    ax.grid(True)

    # 保存或显示
    if save_path:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        plt.savefig(save_path)
        plt.close()
    else:
        plt.show()


def visualize_prediction(true_state, pred_state, action, save_path=None):
    """可视化真实状态和预测状态的比较"""
    # 导入这里避免循环导入
    import matplotlib.pyplot as plt

    action_names = ["UP", "DOWN", "LEFT", "RIGHT", "STAY"]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8))

    visualize_state(true_state, ax1, title=f"真实状态 (动作: {action_names[action]})")
    visualize_state(pred_state, ax2, title="预测状态")

    # 保存或显示
    if save_path:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        plt.savefig(save_path)
        plt.close()
    else:
        plt.show()


def visualize_state(grid, ax, title="State"):
    """在给定轴上可视化状态"""
    colors = {
        EMPTY: 'white',
        START: 'green',
        AGENT: 'blue',
        COIN: 'yellow',
        OBSTACLE: 'black',
        END: 'red'
    }

    names = {
        EMPTY: 'Empty',
        START: 'Start',
        AGENT: 'Agent',
        COIN: 'Coin',
        OBSTACLE: 'Obstacle',
        END: 'End'
    }

    ax.set_title(title)

    # 绘制网格
    height, width = grid.shape
    for i in range(height):
        for j in range(width):
            cell_value = grid[i, j]
            color = colors.get(cell_value, 'gray')
            rect = plt.Rectangle((j, height - i - 1), 1, 1, color=color, ec='black')
            ax.add_patch(rect)

            # 添加文本标签
            if cell_value not in [EMPTY, OBSTACLE]:
                ax.text(j + 0.5, height - i - 0.5, names.get(cell_value, 'Unknown'),
                        ha='center', va='center', fontsize=8)

    # 设置坐标轴
    ax.set_xlim(0, width)
    ax.set_ylim(0, height)
    ax.set_xticks(range(width))
    ax.set_yticks(range(height))
    ax.set_xticklabels(range(width))
    ax.set_yticklabels(range(height - 1, -1, -1))
    ax.grid(True)


def main():
    """主函数"""
    # 设置随机种子
    torch.manual_seed(RANDOM_SEED)
    np.random.seed(RANDOM_SEED)

    # 创建模型目录
    os.makedirs("../models", exist_ok=True)

    # 使用GPU如果可用
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")

    # 1. 生成高级数据集
    print("\n====== 生成高级训练数据集 ======")
    train_data, val_data, test_data = generate_advanced_dataset(num_samples=NUM_SAMPLES)

    # 2. 准备数据集
    print("\n====== 准备数据集 ======")
    train_dataset, val_dataset, test_dataset = prepare_datasets(train_data, val_data, test_data)

    # 3. 创建高级混合世界模型
    print("\n====== 创建高级混合世界模型 ======")
    # 计算节点特征数
    sample_state = train_data[0][0]
    features_per_cell = sample_state.shape[0] // (ENV_SIZE * ENV_SIZE)
    node_features = features_per_cell + 2  # +2 for positional encoding

    print(f"节点特征数: {node_features}")
    print(f"图神经网络隐藏层维度: {GNN_HIDDEN_CHANNELS}")
    print(f"图神经网络层数: {GNN_NUM_LAYERS}")

    model = HybridWorldModel(
        node_features=node_features,
        edge_features=GNN_EDGE_FEATURES,
        hidden_dim=GNN_HIDDEN_CHANNELS,
        history_length=5,
        dropout_rate=DROPOUT_RATE
    ).to(device)

    # 4. 训练模型
    print("\n====== 开始高级模型训练 ======")
    trained_model = advanced_train(model, train_dataset, val_dataset, epochs=NUM_EPOCHS)

    # 5. 保存最终模型
    print("\n====== 保存训练完成的模型 ======")
    torch.save(trained_model.state_dict(), "models/hybrid_world_model_final.pth")

    # 6. 评估模型
    print("\n====== 评估高级模型 ======")
    env_acc, agent_abs_acc, agent_rel_acc = evaluate_advanced_model(trained_model, test_dataset, device)

    # 7. 运行集成预测测试
    print("\n====== 运行集成预测测试 ======")
    overall_accuracy = run_ensemble_episodes(num_episodes=5, steps_per_episode=20)

    print("\n====== 最终结果汇总 ======")
    print(f"测试集环境状态准确率: {env_acc:.4f}")
    print(f"测试集智能体绝对位置准确率: {agent_abs_acc:.4f}")
    print(f"测试集智能体相对移动准确率: {agent_rel_acc:.4f}")
    print(f"集成预测智能体位置准确率: {overall_accuracy:.4f}")


if __name__ == "__main__":
    main()
